import random
import numpy as np
import torch
from torch import nn
import gymnasium as gym


class ReplayBuffer:
    def __init__(self, memory_size):
        self.memory_size = memory_size
        self.memory = []
        self.position = 0

    def record(self, state, action, reward, next_state, done):
        if not self.is_ready():
            self.memory.append([state, action, reward, next_state, done])
            self.position = (self.position + 1) % self.memory_size
        else:
            self.memory[self.position] = [state, action, reward, next_state, done]
            self.position = (self.position + 1) % self.memory_size

    def is_ready(self):
        return len(self.memory) == self.memory_size

    def sample(self, batch_size):
        if not self.is_ready():
            raise ValueError("Not enough samples in memory")
        batch = random.sample(self.memory, batch_size)
        return batch
        # state, action, reward, next_state, done = map(np.stack, zip(*batch))
        # return state, action, reward, next_state, done


class Network(nn.Module):
    """
    input: state->s
    output: qa(s,a1),qa(s,a2),...,qa(s,an)
    """

    def __init__(self, input: int, output: int) -> None:
        super(Network, self).__init__()
        self.fc1 = nn.Linear(input, 64)
        self.fc4 = nn.Linear(64, 16)
        self.fc5 = nn.Linear(16, output)

    def forward(self, state: torch.Tensor):
        state = nn.functional.relu(self.fc1(state))
        state = nn.functional.relu(self.fc4(state))
        state = self.fc5(state)
        return state

    def save(self, path: str):
        torch.save(self.state_dict(), path)

    def load(self, path: str):
        self.load_state_dict(torch.load(path))


class DQN:
    def __init__(
        self,
        state_shape_len: int,
        action_len: int,
        lr: float,
        gamma: float,
        epsilon: float,
        batch_size: int,
        memory_size: int,
        update_freq: int,
    ):
        self.state_shape_len = state_shape_len
        self.action_len = action_len
        self.lr = lr
        self.gamma = gamma
        self.epsilon = epsilon
        self.batch_size = batch_size
        self.update_freq = update_freq
        self.memory = ReplayBuffer(memory_size)
        self.QNet = Network(state_shape_len, action_len)
        self.QTargetNet = Network(state_shape_len, action_len)
        self.QTargetNet.load_state_dict(self.QNet.state_dict())
        self.QNet.train()
        self.QTargetNet.eval()
        self.criterion = torch.nn.MSELoss()
        self.optimizer = torch.optim.Adam(self.QNet.parameters(), lr=self.lr)
        self.freq_count = 0

    def sample_action(self, state):
        if np.random.random() < self.epsilon:
            return np.random.randint(self.action_len)
        else:
            return self.chose_action(state)

    def chose_action(self, state):
        with torch.no_grad():
            return self.QNet(state).argmax().item()

    def state2input(self, state: int):
        input = np.zeros(16)
        input[state] = 1
        return torch.tensor(input, dtype=torch.float32)

    def optimize(self, batch):
        target_list = []
        prediction_list = []
        for state, action, reward, next_state, done in batch:
            if done:
                target = reward
            else:
                with torch.no_grad():
                    target = (
                        reward
                        + self.gamma
                        * self.QTargetNet(self.state2input(next_state)).max()
                    )
            prediction_list.append(self.QNet(self.state2input(state)))
            targetQ = self.QTargetNet(self.state2input(state))
            targetQ[action] = target
            target_list.append(targetQ)
        loss = self.criterion(torch.stack(prediction_list), torch.stack(target_list))
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def train(self, episodes: int, render=False):
        env = gym.make("FrozenLake-v1", render_mode=None, is_slippery=False)
        self.epsilon = 1.0
        self.QTargetNet.load_state_dict(self.QNet.state_dict())
        step_count = 0
        for episode in range(episodes):
            state, info = env.reset()
            done = False
            while not done:
                action = self.sample_action(self.state2input(state))
                next_state, reward, terminated, truncated, info = env.step(action)
                self.memory.record(state, action, reward, next_state, terminated)
                state = next_state
                if terminated or truncated:
                    done = True
            step_count += 1
            self.epsilon = max(self.epsilon - 1 / episodes, 0)
            if self.memory.is_ready():
                batches = self.memory.sample(self.batch_size)
                self.optimize(batches)
            if step_count > self.update_freq:
                self.QTargetNet.load_state_dict(self.QNet.state_dict())
                step_count = 0
            if episode % 10 == 0:
                print(f"episode: {episode}, epsilon: {self.epsilon}")
        env.close()
        torch.save(self.QNet.state_dict(), "dqnnn.pt")

    def test(self):
        env = gym.make("FrozenLake-v1", render_mode="human", is_slippery=False)
        for _ in range(10):
            state, info = env.reset()
            done = False
            while not done:
                action = self.chose_action(self.state2input(state))
                print(f"s:{state},a:{action}")
                state, reward, terminated, truncated, info = env.step(action)
                if terminated or truncated:
                    done = True
        env.close()


if __name__ == "__main__":
    dqn = DQN(
        state_shape_len=16,
        action_len=4,
        gamma=0.9,
        epsilon=1,
        lr=0.001,
        batch_size=64,
        memory_size=2000,
        update_freq=10,
    )
    dqn.train(episodes=2000)
    print("Train Finished")
    dqn.test()
